#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6
import os
import matplotlib
import numpy as np
from torch import device
import time
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import copy
from torchvision import datasets, transforms
import torch
import torch.utils.data as data
from torchvision.datasets import ImageFolder
from utils.sampling import driver_noniid
from utils.options import args_parser
from models.Update import LocalUpdate
from models.Nets import EfficientNetModel
from models.Fed import FedAvg
from sklearn.model_selection import train_test_split
from models.test import test_img
import torch.nn.functional as F


def calculate_sv(models):
    import itertools
    import copy
    """
    计算客户的Shapley值

    Parameters:
    models（dict）：客户端标识符和模型更新的键值对。

    Returns:
    sv: 客户端标识符和计算的shapley值的键值对。
    """
    test_model = EfficientNetModel(args=args)  # 每次聚合后用于测试的模型
    id_list = list(models.keys())
    # 生成可能的排列
    all_perms = list(itertools.permutations(id_list))
    # 所有排列的边际贡献列表
    marginal_contributions = []
    history = {}

    for perm in all_perms:
        perm_values = {}
        local_models = []

        for client_id in perm:
            models_id_index = [i for i, x in enumerate(id_list) if x == client_id]  # 找到该id对应在models列表中的索引位置
            model = copy.deepcopy(models[models_id_index[0]])
            local_models.append(model)
            # 如果当前排列组合的字典 perm_values 是空字典，则将该客户端 ID 转换为一个元素的元组并赋值给 index；否则，将所有键和该客户端 ID 都添加到一个列表中，并按字母顺序排序，
            # 然后将其转换为元组并赋值给 index。
            if len(perm_values.keys()) == 0:
                index = (client_id,)
            else:
                index = tuple(sorted(list(tuple(perm_values.keys()) + (client_id,))))
            # 如果 index 在 history 字典中，则从中获取当前值；
            # 否则，对于 local_models 列表执行 FedAvg 聚合，并将结果加载到模型对象 test_model 中，并执行测试函数 test，将测试结果存储到 history 字典中。
            if index in history.keys():
                current_value = history[index]
            else:
                # 聚合权重
                model_updated_weights = FedAvg(local_models)
                # Load updated weights加载更新的权重
                test_model.load_state_dict(model_updated_weights)
                # 测试准确率
                current_value = test(test_model, test_loader)
                history[index] = current_value

            perm_values[client_id] = max(0, current_value - sum(perm_values.values()))

        marginal_contributions.append(perm_values)

    sv = {client_id: 0 for client_id in id_list}

    # n个客户端每个客户端的边际贡献总和
    for perm in marginal_contributions:
        for key, value in perm.items():
            sv[key] += value

    # 计算平均边际贡献
    sv = {key: value / len(marginal_contributions) for key, value in sv.items()}

    return sv


def shapley_input(dict_user):
    """
    Args:
        dict_user: 客户端
        loaders:划分完成的数据集
    Returns:
    客户端标识符和模型更新的键值对
    """
    # 创建空字典
    input_list = {}
    for idx in dict_user:
        local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
        w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device))
        input_list[idx] = copy.deepcopy(w)
    return input_list


def test(model, testloader):
    # 将模型移动到指定设备（GPU/CPU）上进行计算
    model.to(args.device)
    model.eval()
    test_loss = 0
    correct = 0
    total = len(testloader.dataset)
    # 使用torch.no_grad()上下文管理器进入无梯度计算模式，以减少内存消耗并加快计算速度
    with torch.no_grad():
        # 循环遍历testloader中的所有图像和标签，其中image是输入图像张量，label是相应的标签张量
        for image, label in testloader:
            image, label = image.to(args.device), label.to(args.device)
            # 使用模型对当前的一批图像执行前向传递（推理），并得到输出
            output = model(image)
            # 计算当前批次的测试损失，并将其添加到之前的测试损失中
            test_loss += F.nll_loss(output, label, reduction='sum').item()
            # 找到输出张量中每行的最大值（概率）和对应的索引（类别），并将其保留在预测张量pred中
            pred = output.argmax(dim=1, keepdim=True)
            # 计算当前批次中正确预测的数目，并将其添加到之前正确预测的数量中
            correct += pred.eq(label.view_as(pred)).sum().item()
    # 计算整个测试集上的准确度（正确预测数除以测试总数）
    accuracy = correct / total
    return accuracy


def get_n_largest_keys(dict, n):
    # 选择字典中value值最大的n个key
    sorted_dict = sorted(dict.items(), key=lambda x: x[1], reverse=True)
    n_largest_keys = [x[0] for x in sorted_dict[:n]]
    return n_largest_keys


def reward_clients_by_shapley(shapley_dict, base_learning_rate):
    # 将 Shapley 值从大到小排序并得到客户端 ID 列表
    sorted_clients = sorted(shapley_dict.keys(), key=lambda x: shapley_dict[x], reverse=True)
    # 给予最大奖励的客户端数量
    n_max_reward = int(len(shapley_dict) * 0.4)
    # 给予最大惩罚的客户端数量
    n_max_punish = int(len(shapley_dict) * 0.4)
    lr_dict = {}
    for i, client_id in enumerate(sorted_clients):
        shapley_value = shapley_dict[client_id]
        if i < n_max_reward:
            # 给予最大奖励
            lr =base_learning_rate * (1 + shapley_value)
        elif i >= len(shapley_dict) - n_max_punish:
            # 给予最大惩罚
            lr = base_learning_rate / (1 + shapley_value)
        else:
            # 给予一定程度的惩罚或奖励
            lr = base_learning_rate * (0.9 + shapley_value)
        lr_dict[client_id] = lr
    return lr_dict


if __name__ == '__main__':
    # parse args
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    # load dataset and split users
    if args.dataset == 'FL_dd':
        data_dir = 'data/driving/new_train'
        trans_driving = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
        # 手动划分训练集和测试集
        img_names = os.listdir(data_dir)
        # 遍历 img_names 列表，将每个文件名与 data_dir 组合成完整的文件路径，并以列表形式存储在 img_paths 中
        img_paths = [os.path.join(data_dir, name) for name in img_names]
        # 根据文件名中的标签信息（例如，文件名为“u1_img01.jpg”则其标签为 1），提取出每张图片对应的标签，以列表形式存储在 labels 中
        labels = [int(name.split('_')[0][1:]) for name in img_names]
        # 使用 train_test_split 函数将 img_paths 和 labels 按照 8:2 的比例随机划分为训练集和测试集，
        # 并分别存储在 train_paths、test_paths、train_labels 和 test_labels 四个列表中。其中 random_state=42 表示了随机数生成器的种子，保证每次运行时划分结果一致
        train_paths, test_paths, train_labels, test_labels = train_test_split(img_paths, labels, test_size=0.2,
                                                                              random_state=42)
        # 使用 ImageFolder 类加载数据集。root 参数指定了数据集根目录，transform 参数指定了数据预处理方法
        dataset_train = ImageFolder(root='data/driving/new_train', transform=trans_driving)
        dataset_test = ImageFolder(root='data/driving/newtest', transform=trans_driving)
        # 将加载好的测试集数据集 dataset_test 构建成一个批量迭代器 test_loader。其中 batch_size=32 表示每次迭代返回 32 张图片，shuffle=True 表示每次迭代时对数据进行随机洗牌。
        # 这里将dataset_test 构建成一个批量迭代器 test_loader是为了方便在计算shapley值的时候调用
        test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=32, shuffle=True)
        # 用 driver_noniid 函数将训练集数据集 dataset_train 按照非独立同分布（Non-IID）的方式划分为多个用户，生成一个字典类型的对象 dict_users。args.num_users 表示需要划分的用户数量
        dict_users = driver_noniid(dataset_train, args.num_users)
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape

    # build model
    if args.model == 'en':
        net_glob = EfficientNetModel(args=args).to(args.device)
    else:
        exit('Error: unrecognized model')
    print(net_glob)
    net_glob.train()

    # copy weights
    w_glob = net_glob.state_dict()

    # training
    loss_train = []
    cv_loss, cv_acc = [], []
    val_loss_pre, counter = 0, 0
    net_best = None
    best_loss = None
    val_acc_list, net_list = [], []

    if args.all_clients:
        print("Aggregation over all clients")
        w_locals = [w_glob for i in range(args.num_users)]
    for iter in range(args.epochs):
        # 定义一个空列表loss_locals，用于存储每个客户端的本地损失
        loss_locals = []
        if not args.all_clients:
            w_locals = []
        # 生成客户端标识符和模型更新的键值对
        shapley_input_list = shapley_input(dict_users)
        # 计算每个客户端的shapley值
        T1 = time.time()
        sv = calculate_sv(shapley_input_list)
        T2 = time.time()
        print("客户端和计算的shapley值的键值对:", sv, T2 - T1)
        # 根据shapley值对学习率进行动态激励
        lr_list = reward_clients_by_shapley(sv, args.lr)
        # 选择一部分客户端（数量为frac * num_users）
        m = max(int(args.frac * args.num_users), 1)
        idxs_users = get_n_largest_keys(sv, m)
        for idx in idxs_users:
            # 根据选择的客户端确定激励或惩罚后的学习率
            args.lr = lr_list[idx]
            # 对每个客户端进行本地更新，得到本地模型w和本地损失loss
            local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
            w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device))
            # 并将其添加到w_locals和loss_locals列表中
            if args.all_clients:
                w_locals[idx] = copy.deepcopy(w)
            else:
                w_locals.append(copy.deepcopy(w))
            loss_locals.append(copy.deepcopy(loss))
        # update global weights
        # 使用FedAvg函数对本轮迭代中所有客户端的权重进行聚合
        w_glob = FedAvg(w_locals)

        # copy weight to net_glob
        net_glob.load_state_dict(w_glob)

        # print loss
        loss_avg = sum(loss_locals) / len(loss_locals)
        print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg))
        loss_train.append(loss_avg)

    # plot loss curve
    plt.figure()
    plt.plot(range(len(loss_train)), loss_train)
    plt.ylabel('train_loss')
    plt.savefig('./save/fed_{}_{}_{}_C{}_iid{}.png'.format(args.dataset, args.model, args.epochs, args.frac, args.iid))

    # testing
    net_glob.eval()
    acc_train, loss_train = test_img(net_glob, dataset_train, args)
    acc_test, loss_test = test_img(net_glob, dataset_test, args)
    print("Training accuracy: {:.2f}".format(acc_train))
    print("Testing accuracy: {:.2f}".format(acc_test))
